#!/usr/bin/env python

# noinspection PyUnresolvedReferences
import vtkmodules.vtkInteractionStyle
# noinspection PyUnresolvedReferences
import vtkmodules.vtkRenderingOpenGL2
from vtkmodules.vtkCommonColor import vtkNamedColors
from vtkmodules.vtkFiltersCore import (
    vtkElevationFilter,
    vtkGlyph3D
)
from vtkmodules.vtkFiltersSources import (
    vtkConeSource,
    vtkCubeSource,
    vtkSphereSource
)
from vtkmodules.vtkImagingCore import vtkRTAnalyticSource
from vtkmodules.vtkImagingGeneral import vtkImageGradient
from vtkmodules.vtkInteractionStyle import vtkInteractorStyleTrackballCamera
from vtkmodules.vtkRenderingCore import (
    vtkActor,
    vtkPolyDataMapper,
    vtkRenderWindow,
    vtkRenderWindowInteractor,
    vtkRenderer
)


# This example uses a "glyph table" to change the shape of the 3d glyphs
# according to a scalar value.

# NOTE: The vtkGlyph3D filter doesn't copy over scalars to the glyphs
# generated by a table like this for some reason...

def main():
    colors = vtkNamedColors()

    # The Wavelet Source is nice for generating a test vtkImageData set
    rt = vtkRTAnalyticSource()
    rt.SetWholeExtent(-2, 2, -2, 2, 0, 0)

    # Take the gradient of the only scalar 'RTData' to get a vector attribute
    grad = vtkImageGradient()
    grad.SetDimensionality(3)
    grad.SetInputConnection(rt.GetOutputPort())

    # Elevation just to generate another scalar attribute that varies nicely over the data range
    elev = vtkElevationFilter()
    # Elevation values will range from 0 to 1 between the Low and High Points
    elev.SetLowPoint(-2, -2, 0)
    elev.SetHighPoint(2, 2, 0)
    elev.SetInputConnection(grad.GetOutputPort())

    # Create simple PolyData for glyph table
    cs = vtkCubeSource()
    cs.SetXLength(0.5)
    cs.SetYLength(1)
    cs.SetZLength(2)
    ss = vtkSphereSource()
    ss.SetRadius(0.25)
    cs2 = vtkConeSource()
    cs2.SetRadius(0.25)
    cs2.SetHeight(0.5)

    # Set up the glyph filter
    glyph = vtkGlyph3D()
    glyph.SetInputConnection(elev.GetOutputPort())

    # Here is where we build the glyph table
    # that will be indexed into according to the IndexMode
    glyph.SetSourceConnection(0, cs.GetOutputPort())
    glyph.SetSourceConnection(1, ss.GetOutputPort())
    glyph.SetSourceConnection(2, cs2.GetOutputPort())

    glyph.ScalingOn()
    glyph.SetScaleModeToScaleByScalar()
    glyph.SetVectorModeToUseVector()
    glyph.OrientOn()
    glyph.SetScaleFactor(1)  # Overall scaling factor
    glyph.SetRange(0, 1)  # Default is (0,1)

    # Tell it to index into the glyph table according to scalars
    glyph.SetIndexModeToScalar()

    # Tell glyph which attribute arrays to use for what
    glyph.SetInputArrayToProcess(0, 0, 0, 0, 'Elevation')  # scalars
    glyph.SetInputArrayToProcess(1, 0, 0, 0, 'RTDataGradient')  # vectors

    coloring_by = 'Elevation'
    mapper = vtkPolyDataMapper()
    mapper.SetInputConnection(glyph.GetOutputPort())
    mapper.SetScalarModeToUsePointFieldData()
    mapper.SetColorModeToMapScalars()
    mapper.ScalarVisibilityOn()

    # GetRange() call doesn't work because attributes weren't copied to glyphs
    # as they should have been...
    # mapper.SetScalarRange(glyph.GetOutputDataObject(0).GetPointData().GetArray(coloring_by).GetRange())

    mapper.SelectColorArray(coloring_by)
    actor = vtkActor()
    actor.SetMapper(mapper)

    ren = vtkRenderer()
    ren.AddActor(actor)
    ren.SetBackground(colors.GetColor3d("DarkGray"))

    renWin = vtkRenderWindow()
    renWin.AddRenderer(ren)
    renWin.SetWindowName('GlyphTable')

    iren = vtkRenderWindowInteractor()
    istyle = vtkInteractorStyleTrackballCamera()
    iren.SetInteractorStyle(istyle)
    iren.SetRenderWindow(renWin)
    ren.ResetCamera()
    renWin.Render()
    iren.Start()


if __name__ == '__main__':
    main()
